#!/usr/bin/python 
# -*- coding: utf8 -*- 
import os
import sys
import boto.ec2
import yaml
from pymongo import MongoClient
from multiprocessing import Pool
from multiprocessing.dummy import Pool as ThreadPool


cur_dir = os.path.dirname(os.path.abspath(__file__))

class MongoSession():
    def __init__(self):
        setting = yaml.load(open(os.path.join(cur_dir,'setting.yaml')))
        self.conn = MongoClient(setting['mongodb'],setting['mongodb_port'])
        self.db = self.conn[setting['db']]
    
    def insertEc2InstanceInfo(self,instance):
        self.db['ec2instance'].insert(instance)
        return None
    
    def removeEc2RegionAllInstance(self,region):
        self.db['ec2instance'].remove({"region":region,"type":"aws"})
        return None

def search_account(name):
    setting = yaml.load(open(os.path.join(cur_dir,'setting.yaml')))
    for acc in setting['account']:
        if acc['name'] == name:
            return acc
    return False

def get_conn_aws(region,key_id,access_key):
    conn = boto.ec2.connect_to_region(region,aws_access_key_id=key_id,aws_secret_access_key=access_key)
    return conn

def check_tags_value(key,tags):
    if key not in tags:
        return None
    else:
        return tags[key]

def instance_dict(inst):
    '''
    一个实例的信息
    '''
    inst_dict = {}
    inst_dict['id'] = inst.id
    inst_dict['type'] = 'aws'
    tags = inst.tags
    inst_dict['department'] = check_tags_value('department',tags)
    inst_dict['owners'] = check_tags_value('owners',tags)
    inst_dict['product'] = check_tags_value('product',tags)
    inst_dict['Name'] = check_tags_value('Name',tags)
    inst_dict['use'] = check_tags_value('use',tags)
    inst_dict['instance_type'] = inst.instance_type
    inst_dict['region'] = inst.region.name
    inst_dict['placement'] = inst.placement
    inst_dict['state'] = inst.state
    inst_dict['public_dns_name'] = inst.public_dns_name
    inst_dict['ip_address'] = inst.ip_address
    inst_dict['launch_time'] = inst.launch_time
    inst_dict['private_dns_name'] = inst.private_dns_name
    inst_dict['private_ip_address'] = inst.private_ip_address
    inst_dict['vpc_id'] = inst.vpc_id
    inst_dict['image_id'] = inst.image_id
    inst_dict['subnet_id'] = inst.subnet_id
    inst_dict['platform'] = inst.platform
    inst_dict['ebs_optimized'] = inst.ebs_optimized
    inst_dict['root_device_name'] = inst.root_device_name
    inst_dict['root_device_type'] = inst.root_device_type
    inst_dict['block_devices'] = [i for i in inst.block_device_mapping ]
    inst_dict['kernel'] = inst.kernel
    return inst_dict
    
def insert_all_instance_info(conn,region):
    mongo = MongoSession()
    reservations = conn.get_all_reservations()
    if reservations:
        #把之前的记录删除
        mongo.removeEc2RegionAllInstance(region)
        for i in reservations:
            _instances = i.instances
            for inst in _instances:
                inst_dict = instance_dict(inst)
                mongo.insertEc2InstanceInfo(inst_dict)
        return True
    else:
        return False

def work(region):
    region_id = region['id']
    readonlyaccount = search_account('readonly')
    access_key_id = readonlyaccount['access_key_id']
    secret_access_key = readonlyaccount['secret_access_key']        
    conn = get_conn_aws(region_id,access_key_id,secret_access_key)
    if conn:
        result = insert_all_instance_info(conn,region_id)
        if result:
            return 'Region:%s insert aws info ok' % region_id
        else:
            return 'Region:%s insert aws info False' % region_id
    else:
        return 'Region:%s conn aws False' % region_id
    
def main():
    setting = yaml.load(open(os.path.join(cur_dir,'setting.yaml')))
    aws_region_list = setting['aws_region']
    pool = ThreadPool(4)
    result = pool.map(work,aws_region_list)
    pool.close()
    pool.join()
    print result
    
if __name__ ==  "__main__":
    main()